import numpy as np
import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk, ImageOps  # 修改后的导入
import time
from plyfile import PlyData, PlyElement
import matplotlib.pyplot as plt
import cv2

color_map = [
    (128, 0, 0),     (132, 0, 0),     (136, 0, 0),     (140, 0, 0),
    (144, 0, 0),     (148, 0, 0),     (152, 0, 0),     (156, 0, 0),
    (160, 0, 0),     (164, 0, 0),     (168, 0, 0),     (172, 0, 0),
    (176, 0, 0),     (180, 0, 0),     (184, 0, 0),     (188, 0, 0),
    (192, 0, 0),     (196, 0, 0),     (200, 0, 0),     (204, 0, 0),
    (208, 0, 0),     (212, 0, 0),     (216, 0, 0),     (220, 0, 0),
    (224, 0, 0),     (228, 0, 0),     (232, 0, 0),     (236, 0, 0),
    (240, 0, 0),     (244, 0, 0),     (248, 0, 0),     (252, 0, 0),
    (255, 0, 0),     (255, 4, 0),     (255, 8, 0),     (255, 12, 0),
    (255, 16, 0),    (255, 20, 0),    (255, 24, 0),    (255, 28, 0),
    (255, 32, 0),    (255, 36, 0),    (255, 40, 0),    (255, 44, 0),
    (255, 48, 0),    (255, 52, 0),    (255, 56, 0),    (255, 60, 0),
    (255, 64, 0),    (255, 68, 0),    (255, 72, 0),    (255, 76, 0),
    (255, 80, 0),    (255, 84, 0),    (255, 88, 0),    (255, 92, 0),
    (255, 96, 0),    (255, 100, 0),   (255, 104, 0),   (255, 108, 0),
    (255, 112, 0),   (255, 116, 0),   (255, 120, 0),   (255, 124, 0),
    (255, 128, 0),   (255, 132, 0),   (255, 136, 0),   (255, 140, 0),
    (255, 144, 0),   (255, 148, 0),   (255, 152, 0),   (255, 156, 0),
    (255, 160, 0),   (255, 164, 0),   (255, 168, 0),   (255, 172, 0),
    (255, 176, 0),   (255, 180, 0),   (255, 184, 0),   (255, 188, 0),
    (255, 192, 0),   (255, 196, 0),   (255, 200, 0),   (255, 204, 0),
    (255, 208, 0),   (255, 212, 0),   (255, 216, 0),   (255, 220, 0),
    (255, 224, 0),   (255, 228, 0),   (255, 232, 0),   (255, 236, 0),
    (255, 240, 0),   (255, 244, 0),   (255, 248, 0),   (255, 252, 0),
    (254, 255, 1),   (250, 255, 6),   (246, 255, 10),  (242, 255, 14),
    (238, 255, 18),  (234, 255, 22),  (230, 255, 26),  (226, 255, 30),
    (222, 255, 34),  (218, 255, 38),  (214, 255, 42),  (210, 255, 46),
    (206, 255, 50),  (202, 255, 54),  (198, 255, 58),  (194, 255, 62),
    (190, 255, 66),  (186, 255, 70),  (182, 255, 74),  (178, 255, 78),
    (174, 255, 82),  (170, 255, 86),  (166, 255, 90),  (162, 255, 94),
    (158, 255, 98),  (154, 255, 102), (150, 255, 106), (146, 255, 110),
    (142, 255, 114), (138, 255, 118), (134, 255, 122), (130, 255, 126),
    (126, 255, 130), (122, 255, 134), (118, 255, 138), (114, 255, 142),
    (110, 255, 146), (106, 255, 150), (102, 255, 154), (98, 255, 158),
    (94, 255, 162),  (90, 255, 166),  (86, 255, 170),  (82, 255, 174),
    (78, 255, 178),  (74, 255, 182),  (70, 255, 186),  (66, 255, 190),
    (62, 255, 194),  (58, 255, 198),  (54, 255, 202),  (50, 255, 206),
    (46, 255, 210),  (42, 255, 214),  (38, 255, 218),  (34, 255, 222),
    (30, 255, 226),  (26, 255, 230),  (22, 255, 234),  (18, 255, 238),
    (14, 255, 242),  (10, 255, 246),  (6, 255, 250),   (2, 255, 254),
    (0, 252, 255),   (0, 248, 255),   (0, 244, 255),   (0, 240, 255),
    (0, 236, 255),   (0, 232, 255),   (0, 228, 255),   (0, 224, 255),
    (0, 220, 255),   (0, 216, 255),   (0, 212, 255),   (0, 208, 255),
    (0, 204, 255),   (0, 200, 255),   (0, 196, 255),   (0, 192, 255),
    (0, 188, 255),   (0, 184, 255),   (0, 180, 255),   (0, 176, 255),
    (0, 172, 255),   (0, 168, 255),   (0, 164, 255),   (0, 160, 255),
    (0, 156, 255),   (0, 152, 255),   (0, 148, 255),   (0, 144, 255),
    (0, 140, 255),   (0, 136, 255),   (0, 132, 255),   (0, 128, 255),
    (0, 124, 255),   (0, 120, 255),   (0, 116, 255),   (0, 112, 255),
    (0, 108, 255),   (0, 104, 255),   (0, 100, 255),   (0, 96, 255),
    (0, 92, 255),    (0, 88, 255),    (0, 84, 255),    (0, 80, 255),
    (0, 76, 255),    (0, 72, 255),    (0, 68, 255),    (0, 64, 255),
    (0, 60, 255),    (0, 56, 255),    (0, 52, 255),    (0, 48, 255),
    (0, 44, 255),    (0, 40, 255),    (0, 36, 255),    (0, 32, 255),
    (0, 28, 255),    (0, 24, 255),    (0, 20, 255),    (0, 16, 255),
    (0, 12, 255),    (0, 8, 255),     (0, 4, 255),     (0, 0, 255),
    (0, 0, 252),     (0, 0, 248),     (0, 0, 244),     (0, 0, 240),
    (0, 0, 236),     (0, 0, 232),     (0, 0, 228),     (0, 0, 224),
    (0, 0, 220),     (0, 0, 216),     (0, 0, 212),     (0, 0, 208),
    (0, 0, 204),     (0, 0, 200),     (0, 0, 196),     (0, 0, 192),
    (0, 0, 188),     (0, 0, 184),     (0, 0, 180),     (0, 0, 176),
    (0, 0, 172),     (0, 0, 168),     (0, 0, 164),     (0, 0, 160),
    (0, 0, 156),     (0, 0, 152),     (0, 0, 148),     (0, 0, 144),
    (0, 0, 140),     (0, 0, 136),     (0, 0, 132),     (0, 0, 128)
]

fx = 2.265142e+02  # 需替换为A010实际焦距
fy = 2.278584e+02
cx = 50.0       # 100x100分辨率下光心位置 (原320x240的163.7246/2)
cy = 50.0      # (原240x320的123.3738/2)

global drop_frame
drop_frame = 0

def visualize_ply(file_path):
    # 读取PLY文件
    ply_data = PlyData.read(file_path)
    vertices = ply_data['vertex']
    
    # 提取坐标数据
    x = vertices['x']
    y = vertices['y']
    z = vertices['z']
    
    # 计算深度范围
    min_depth = np.min(z)
    max_depth = np.max(z)
    depth_range = max_depth - min_depth
    
    # 生成颜色映射
    colors = []
    for depth in z:
        # 将深度映射到颜色索引 (0-255)
        color_idx = int(((depth - min_depth) / depth_range) * (len(color_map)-1))
        colors.append(color_map[color_idx])
    
    # 转换为0-1范围的RGB值
    colors = np.array(colors) / 255.0

    # 创建3D可视化
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制点云
    ax.scatter(x, y, z, c=colors, s=1, marker='o')
    
    # 设置坐标轴标签
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    
    # 启用交互式旋转
    plt.show(block=True)  # 添加block=True参数保持窗口显示

def generate_point_clouds(filename,frames):
    """生成PLY格式点云文件，同一批次使用相同时间戳"""
    savename=filename.split(".")[0]
    if not savename:
        savename=filename
    
    for frame_idx, frame_data in enumerate(frames):
        depth_data = frame_data['depth']
        depths = []
        point_data = frame_data['point_cloud']  # 改用三维坐标数据
        points = []
        height, width = depth_data.shape
        
        # 生成点坐标
        for y in range(height):
            for x in range(width):
                x_real, y_real, z = point_data[y, x]
                if z>0:
                    points.append((x_real, y_real, z))  # 使用真实坐标
                if (z := depth_data[y, x]) > 0:
                    depths.append((x, y, z))  # 使用元组格式
        
        if depths:
            # 创建PLY数据结构
            vertex = np.array(depths, dtype=[
                ('x', 'f4'), 
                ('y', 'f4'),
                ('z', 'f4')
            ])
            
            el = PlyElement.describe(vertex, 'vertex')            
            PlyData([el]).write(f"./depth/{savename}_frame{frame_idx}.ply")

        if points:
            # 创建PLY数据结构
            vertex = np.array(points, dtype=[
                ('x', 'f4'), 
                ('y', 'f4'),
                ('z', 'f4')
            ])
            
            el = PlyElement.describe(vertex, 'vertex')            
            PlyData([el]).write(f"./point/{savename}_frame{frame_idx}.ply")

def read_sensor_data(path):
    """读取并处理传感器原始数据"""
    try:
        with open(path, "r") as f:
            return [int(x, 16) if x else 0 for x in f.read().split(" ")]
    except:
        pass

def read_binary_data(path):
    """读取二进制文件数据"""
    with open(path, "rb") as f:
        return list(f.read())

def process_frames(raw_data,unit=1):
    global drop_frame
    drop_frame = 0

    """解析原始数据生成图像帧序列"""
    frame_index = 0
    image_array = np.zeros((100, 100, 3), dtype=np.uint8)
    depth_array = np.zeros((100, 100), dtype=np.float32)  # 新增深度数据存储
    point_cloud = np.zeros((100, 100, 3), dtype=np.float32) 
    frame_list = []

    while frame_index < len(raw_data):
        if (frame_index + 3 < len(raw_data) and
            raw_data[frame_index] == 0 and
            raw_data[frame_index+1] == 255 and
            raw_data[frame_index+2] == 32 and
            raw_data[frame_index+3] == 39):
            
            frame_start = frame_index
            frame_index += 20  # 跳过帧头(4) + 保留数据(16)
            
            if frame_index + 10000 + 2 > len(raw_data):
                drop_frame += 1
                frame_index = frame_start + 1
                continue
                
            checksum1 = raw_data[frame_index + 10000]
            checksum2 = raw_data[frame_index + 10000 + 1]
            calculated = sum(raw_data[frame_start:frame_index+10000]) % 256
            
            if checksum1 != calculated or checksum2 != 221:
                drop_frame += 1
                frame_index = frame_start + 1
                # print(frame_index,checksum2)
                continue

            for i in range(10000):
                    y, x = i // 100, i % 100
                    depth = raw_data[frame_index + i]
                    # 直接使用RGB格式 (移除[::-1]的BGR转换)
                    image_array[y, x] = color_map[depth % len(color_map)]

                    if depth == 255:  # 原始数据最大值
                        z = -1.0
                    else:
                        if unit:
                            z = depth * unit  # 新增深度计算
                        else:
                            z = depth*depth/26.01

                    x_real = (x - cx) * z / fx
                    y_real = (y - cy) * z / fy

                    depth_array[y, x] = z
                    point_cloud[y, x] = [x_real, y_real, z]
            
            frame_list.append({
                'color': image_array.copy(),
                'depth': depth_array.copy(),
                'point_cloud': point_cloud.copy()  # 新增点云数据存储
            })
            frame_index = frame_start + 4 + 16 + 10000 + 2
        else:
            frame_index += 1
    return frame_list

class AnimationApp:
    def __init__(self, master):
        self.master = master
        master.title("Depth Sensor Viewer")
        master.geometry("1000x1000")
        
        self.control_frame = tk.Frame(master)
        self.control_frame.pack(pady=10)

        # 新增信息显示面板
        self.info_frame = tk.Frame(master)
        self.info_frame.pack(pady=5)
        
        # 当前帧数显示
        tk.Label(self.info_frame, text="当前帧号:").pack(side=tk.LEFT)
        self.current_frame_entry = tk.Entry(self.info_frame, width=10, state='readonly')
        self.current_frame_entry.pack(side=tk.LEFT, padx=5)
        
        # 实时FPS显示
        tk.Label(self.info_frame, text="实时FPS:").pack(side=tk.LEFT)
        self.real_fps_entry = tk.Entry(self.info_frame, width=10, state='readonly')
        self.real_fps_entry.pack(side=tk.LEFT, padx=5)
        
        # 初始化FPS计算相关变量
        self.frame_times = []
        self.last_time = 0

        # FPS调节器
        self.fps_var = tk.IntVar(value=15)
        tk.Label(self.control_frame, text="FPS:").pack(side=tk.LEFT)
        self.spinbox = tk.Spinbox(self.control_frame, from_=1, to=19, width=5, 
                                textvariable=self.fps_var, state='readonly')  # 添加只读状态
        self.spinbox.pack(side=tk.LEFT, padx=5)

        # 新增UNIT设置
        self.unit_var = tk.IntVar(value=1)
        tk.Label(self.control_frame, text="UNIT:").pack(side=tk.LEFT)
        self.unit_spinbox = tk.Spinbox(self.control_frame, from_=0, to=9, width=5,
                                     textvariable=self.unit_var, state='readonly')
        self.unit_spinbox.pack(side=tk.LEFT, padx=5)

        # 新增播放控制按钮
        self.play_btn = tk.Button(self.control_frame, text="播放", width=8, command=self.start_animation)
        self.play_btn.pack(side=tk.LEFT, padx=5)
        self.stop_btn = tk.Button(self.control_frame, text="暂停", width=8, command=self.stop_animation)
        self.stop_btn.pack(side=tk.LEFT, padx=5)
        self.stop_btn = tk.Button(self.control_frame, text="停止", width=8, command=self.clear_animation)
        self.stop_btn.pack(side=tk.LEFT, padx=5)

        # 新增点云显示按钮
        self.pc_btn = tk.Button(self.control_frame, text="显示深度点云", width=12, command=self.show_depth)
        self.pc_btn.pack(side=tk.LEFT, padx=5)

        self.pc_btn = tk.Button(self.control_frame, text="显示3D点云", width=12, command=self.show_point)
        self.pc_btn.pack(side=tk.LEFT, padx=5)

        self.pc_btn = tk.Button(self.control_frame, text="保存视频", width=12, command=self.save_video)
        self.pc_btn.pack(side=tk.LEFT, padx=5)

        # 创建顶部控制面板
        self.frame = tk.Frame(master)
        self.frame.pack(pady=10)
        
        # 文件路径输入框
        self.entry = tk.Entry(self.frame, width=50, state='readonly')
        self.entry.pack(side=tk.LEFT, padx=5)
        
        # 新增帧数显示标签
        self.frame_count_label = tk.Label(self.frame, width=25, text="已解析帧数: 0")
        self.frame_count_label.pack(side=tk.LEFT, padx=10)

        # 在文本文件控制面板添加丢帧标签
        self.drop_label = tk.Label(self.frame, width=25, text="丢帧: 0")
        self.drop_label.pack(side=tk.LEFT, padx=10)
        
        # 文件选择按钮
        self.btn = tk.Button(self.frame, width=25,text="选择文本文件", command=self.load_file)
        self.btn.pack(side=tk.LEFT)

        # 新增二进制文件控制面板
        self.binary_frame = tk.Frame(master)
        self.binary_frame.pack(pady=10)
        
        # 二进制文件输入框
        self.binary_entry = tk.Entry(self.binary_frame, width=50, state='readonly')
        self.binary_entry.pack(side=tk.LEFT, padx=5)
        
        # 二进制文件帧数标签
        self.binary_frame_label = tk.Label(self.binary_frame, width=25, text="已解析帧数: 0")
        self.binary_frame_label.pack(side=tk.LEFT, padx=10)

        # 在二进制文件控制面板添加丢帧标签
        self.binary_drop_label = tk.Label(self.binary_frame, width=25, text="丢帧: 0")
        self.binary_drop_label.pack(side=tk.LEFT, padx=10)

        # 二进制文件选择按钮
        self.binary_btn = tk.Button(self.binary_frame, width=25,text="选择二进制文件", command=self.load_binary_file)
        self.binary_btn.pack(side=tk.LEFT)
        
        # 图像显示区域
        self.panel = tk.Label(master)
        self.panel.pack(expand=True)
        
        # 动画控制变量
        self.is_playing = False
        self.current_frames = []
        self.frame_index = 0

    def save_video(self):
        self.stop_animation()  # 暂停播放
        self.frame_index = 0  # 重置当前帧号

        self.current_frame_entry.config(state='normal')
        self.current_frame_entry.delete(0, tk.END)
        self.current_frame_entry.insert(0, str(self.frame_index + 1))
        self.current_frame_entry.config(state='readonly')

        if not self.current_frames or not hasattr(self, 'current_filename'):
            return
    
        # 创建视频目录
        save_path = f"./video/{self.current_filename}_video.mp4"
        
        # 获取视频参数
        fps = self.fps_var.get()
        frame_size = (800, 800)
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        
        # 创建视频写入器
        out = cv2.VideoWriter(save_path, fourcc, fps, frame_size)
        
        # 遍历所有帧（从第0帧到最后一帧）
        for idx in range(len(self.current_frames)):
            frame_data = self.current_frames[idx]
            rgb_frame = frame_data['color']
            
            # 转换颜色空间并调整尺寸
            bgr_frame = cv2.cvtColor(rgb_frame, cv2.COLOR_RGB2BGR)
            resized_frame = cv2.resize(bgr_frame, frame_size, interpolation=cv2.INTER_NEAREST)
            out.write(resized_frame)
        
        out.release()

    def show_point(self):
        """显示当前帧点云"""
        self.stop_animation()  # 先暂停播放
        
        if not self.current_frames or not hasattr(self, 'current_filename'):
            return
            
        frame_number = self.frame_index
        ply_path = f"./point/{self.current_filename}_frame{frame_number}.ply"
        
        if os.path.exists(ply_path):
            visualize_ply(ply_path)

    def show_depth(self):
        """显示当前帧点云"""
        self.stop_animation()  # 先暂停播放
        
        if not self.current_frames or not hasattr(self, 'current_filename'):
            return
            
        frame_number = self.frame_index
        ply_path = f"./depth/{self.current_filename}_frame{frame_number}.ply"
        
        if os.path.exists(ply_path):
            visualize_ply(ply_path)

    def load_data(self,flag):
        global drop_frame
        self.clear_animation()
        drop_frame = 0  # 解析前清零
        
        filepath = filedialog.askopenfilename()
        if not filepath:
            return
        
        self.current_filename = os.path.basename(filepath).split('.')[0]

        if flag:            
            self.binary_entry.config(state='normal')
            self.binary_entry.delete(0, tk.END)
            self.binary_entry.insert(0, filepath)
            self.binary_entry.config(state='readonly')

            data = read_binary_data(filepath)
        else:
             # 更新输入框内容（需要临时启用写入）
            self.entry.config(state='normal')
            self.entry.delete(0, tk.END)
            self.entry.insert(0, filepath)
            self.entry.config(state='readonly')
            
            # 处理新数据
            data = read_sensor_data(filepath)
        
        # 读取二进制数据        
        filename = os.path.basename(filepath)
        self.current_frames = process_frames(data, self.unit_var.get())

        generate_point_clouds(filename,self.current_frames)  # 调用全局函数

        if flag:  
            self.binary_frame_label.config(text=f"已解析帧数: {len(self.current_frames)}")
            self.binary_drop_label.config(text=f"丢帧: {drop_frame}")  # 更新显示
        else:
            self.frame_count_label.config(text=f"已解析帧数: {len(self.current_frames)}")
            self.drop_label.config(text=f"丢帧: {drop_frame}")  # 更新显示

        self.start_animation()

    def load_binary_file(self):
        self.load_data(1)

    def load_file(self):
        self.load_data(0)


    def start_animation(self):
        """启动动画循环"""
        if not self.current_frames:
            return
        
        self.is_playing = True
        self.update_frame()

    def stop_animation(self):
        """停止动画"""
        self.is_playing = False
        # self.current_frames = []
        # self.frame_index = 0
        # global drop_frame
        # drop_frame = 0  # 解析前清零
    
    def clear_animation(self):
        """清空动画"""
        self.is_playing = False
        self.current_frames = []
        self.frame_index = 0
        global drop_frame
        drop_frame = 0  # 解析前清零

    def update_frame(self):
        """更新显示帧"""        
        if not self.is_playing:
            return
            
        try:
            frame_data = self.current_frames[self.frame_index]  # 获取包含颜色和深度的帧数据
            color_frame = frame_data['color']  # 提取颜色数据
        except:
            if not self.is_playing:
                return
        
        # 更新当前帧号显示
        self.current_frame_entry.config(state='normal')
        self.current_frame_entry.delete(0, tk.END)
        self.current_frame_entry.insert(0, str(self.frame_index + 1))
        self.current_frame_entry.config(state='readonly')

        # 计算实时FPS        
        current_time = time.time()
        if self.last_time != 0:
            frame_interval = current_time - self.last_time
            self.frame_times.append(frame_interval)
            if len(self.frame_times) > 5:  # 取最近5帧计算平均FPS
                self.frame_times.pop(0)
            real_fps = 1 / (sum(self.frame_times) / len(self.frame_times))
            
            # 更新FPS显示
            self.real_fps_entry.config(state='normal')
            self.real_fps_entry.delete(0, tk.END)
            self.real_fps_entry.insert(0, f"{real_fps:.1f}")
            self.real_fps_entry.config(state='readonly')

        self.last_time = current_time

        # 替换OpenCV的resize和颜色转换
        pil_img = Image.fromarray(color_frame)
        resized = pil_img.resize((800, 800), Image.NEAREST)  # 使用PIL的最近邻缩放
        
        imgtk = ImageTk.PhotoImage(image=resized)
        
        self.panel.config(image=imgtk)
        self.panel.image = imgtk

        self.frame_index = (self.frame_index + 1) % len(self.current_frames)
        target_interval = 1000 / self.fps_var.get()  # 目标间隔（毫秒）
        elapsed = (time.time() - self.last_time) * 1000  # 实际耗时（毫秒）

        remaining_delay = max(0, int(target_interval - elapsed))     

        if self.is_playing:    
            self.master.after(remaining_delay, self.update_frame)

        

if __name__ == "__main__":
    import os
    os.makedirs('./point', exist_ok=True)  # 新增目录检查
    os.makedirs('./depth', exist_ok=True)  # 新增目录检查
    os.makedirs('./video', exist_ok=True)  # 新增目录检查
    
    root = tk.Tk()
    app = AnimationApp(root)
    root.mainloop()
